package nl.cypherpunk.proLearner.smtp;

import java.io.*;
import java.net.*;
import java.nio.ByteBuffer;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.text.DecimalFormat;
import java.util.*;

import static nl.cypherpunk.proLearner.Common.intToByteArray;

public class SMTPTestService {
    Socket socket;
    OutputStream output;
    InputStream input;

    String host = "127.0.0.1";
    int port = 25;
    String hostname = "localhost";
    // Act as a TLS client
    boolean ROLE_CLIENT = true;
    // Restart server after every session
    boolean REQUIRE_RESTART = true;
    // Timeout in ms
    int RECEIVE_MSG_TIMEOUT = 100;
    // Send output from TLS implementation to console
    boolean CONSOLE_OUTPUT = false;

    String cmd;
    public Process targetProcess;

    String redirectFile;
    ArrayList<byte[]> currentCase;
    int totalExecs;
    int prevCoverage;
    String outputDir;
    String env_cmd="";
    PrintWriter out;
    BitSet coverageMap;
    Double initScore;
    String SendUser="ubuntu@ubuntu";
    String SendPassword="xlj932834897";
    String ReceiveUser="ubuntu@ubuntu";
    MappedByteBuffer fileCoverage;
    HashSet<String> crashHash;
    String version = "";
    ArrayList<String> symbols;
    Map<String,Integer> subtypes;
    LinkedHashMap<String, byte[]> mapper;
    HashSet<ByteBuffer> messages;

    public SMTPTestService() throws Exception {
        symbols = new ArrayList<>();
        messages = new HashSet<>();
        subtypes = new LinkedHashMap<>();
        mapper = new LinkedHashMap<>();
    }

    public void setTarget(String target) throws Exception {
        if(target.equals("server")) {
            ROLE_CLIENT = true;
        }
        else if(target.equals("client")) {
            ROLE_CLIENT = false;
        }
        else {
            throw new Exception("Unknown target");
        }
    }

    public void setHost(String host) {
        this.host = host;
    }

    public void setPort(int port) {
        this.port = port;
    }

    public void setCommand(String cmd) {
        this.cmd = cmd;
    }

    public void setRestartTarget(boolean restart) {
        this.REQUIRE_RESTART = restart;
    }

    public void setReceiveMessagesTimeout(int timeout) {
        RECEIVE_MSG_TIMEOUT = timeout;
    }

    public void setConsoleOutput(boolean enable) {
        CONSOLE_OUTPUT = enable;
    }

    public void setOutputDir(String dir){
        this.outputDir = dir;
    }
    public void setRedirectFile(String output) {
        redirectFile = output;
    }
    public void setVersion(String version) {
        this.version = version;
    }
    //update mapper and subtypes
    public boolean addMapper(String type, byte[] message){
        String symbol="";
        if(messages.contains(ByteBuffer.wrap(message)))
            return false;
        messages.add(ByteBuffer.wrap(message));
        if(subtypes.containsKey(type))
        {
            Integer index = subtypes.get(type)+1;
            subtypes.put(type,index);
            symbol=type+(new DecimalFormat("00").format(index));
        }
        else
        {
            subtypes.put(type,1);
            symbol = type+"01";
        }
        symbols.add(symbol);
        mapper.put(symbol,message);
        return true;
    }
    //reverse order
    public boolean delMapper(String subtype){
        symbols.remove(subtype);
        byte[] tmp = mapper.remove(subtype);
        messages.remove(ByteBuffer.wrap(tmp));
        String type = subtype.substring(0,subtype.length()-2);
        if(subtypes.get(type)==1)
        {
            subtypes.remove(type);
        }
        else
        {
            subtypes.put(type,subtypes.get(type)-1);
        }
        return true;
    }

    public ArrayList<String> getSymbols()
    {
        return symbols;
    }

    public void saveMapper()
    {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        for(Map.Entry<String, byte[]> entry : mapper.entrySet()) {
            try {
                file = new File(this.outputDir + File.separator + entry.getKey());
                file.createNewFile();
                fos = new FileOutputStream(file);
                bos = new BufferedOutputStream(fos);
                byte[] tmp = intToByteArray(entry.getValue().length);
                bos.write(tmp, 0, 4);
                bos.write(entry.getValue(), 0, entry.getValue().length);

            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (bos != null) {
                    try {
                        bos.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
                if (fos != null) {
                    try {
                        fos.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
            }
        }
    }

    public void start() throws Exception {
        currentCase = new ArrayList<>();
        prevCoverage = 0;
        totalExecs = 0;
        crashHash = new HashSet<>();
        if(ROLE_CLIENT) {
            /*String killall = "sudo killall exim";
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();*/
            String killall = "sudo rm /home/exim"+version;
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();
            if(cmd != null && !cmd.equals("")) {
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //pb.redirectOutput(new File(redirectFile));
                targetProcess = pb.start();
                Thread.sleep(10000);
            }
            //connectSocket();
        }
        else {
        }
        fileCoverage = new RandomAccessFile("/home/exim"+version, "r").getChannel().map(FileChannel.MapMode.READ_ONLY, 0, 40008);
    }

    public void reset() throws Exception {
        //System.out.println("RESET");
        if(socket!=null)
        {
            socket.close();
        }
        totalExecs ++;
        if(totalExecs%1000==0)
            System.out.println(totalExecs);
        if(ROLE_CLIENT) {
            if (!currentCase.isEmpty()) {
                currentCase.clear();
            }
            if (REQUIRE_RESTART && cmd != null && !cmd.equals("")) {
                targetProcess.destroy();

                String target = this.outputDir.substring(0, this.outputDir.indexOf("2"));
                String killall = "sudo killall " + target;
                if (REQUIRE_RESTART) {
                    ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                    Process tmp = pbkill.start();
                    tmp.waitFor();
                }
                Thread.sleep(100);

                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);

                targetProcess = pb.start();

                Thread.sleep(5000);
            }
            Thread.sleep(100);
            String tmp = connectSocket();
            int i = 0;
            while(!tmp.equals("220"))
            {
                targetProcess.destroy();

                String target = this.outputDir.substring(0, this.outputDir.indexOf("2"));
                String killall = "sudo killall " + target;

                ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                Process kill = pbkill.start();
                kill.waitFor();

                Thread.sleep(100);

                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);

                targetProcess = pb.start();

                Thread.sleep(5000);
                tmp = connectSocket();
            }

        } else {

        }
    }

    public String connectSocket() throws UnknownHostException, IOException,Exception {
        socket = new Socket(host, port);
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);
        output = socket.getOutputStream();
        input = socket.getInputStream();
        return receiveMessages();
    }

    public void listenSocket() throws UnknownHostException, IOException {
        ServerSocket server = new ServerSocket();
        server.bind(new InetSocketAddress(host, port));
        socket = server.accept();
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);

        output = socket.getOutputStream();
        input = socket.getInputStream();

        server.close();
    }

    public void closeSocket() throws IOException, Exception {
        socket.close();
    }

    public String receiveMessages() throws Exception {
        String out = "";
        String statusCode = null;
        if(socket.isInputShutdown())
            return "ConnectionClosed";
        try {
            byte[] code = new byte[3];
            input.read(code, 0, 3);
            StringBuffer tStringBuf=new StringBuffer ();
            char[] tChars=new char[3];

            for(int i=0;i<3;i++)
                tChars[i]=(char)code[i];

            tStringBuf.append(tChars);
            out=tStringBuf.toString();

            if(tChars[0]< 0x31 || tChars[0]> 0x40)
                out = "ConnectionClosed";
            if(out.equals("221"))
                out = "221ConnectionClosed";
            if(out.equals("500"))
                out = "554";

        }
        catch (SocketTimeoutException e)
        {
            //socket.close();
            return "Empty";
        }
        /*
        byte[] data = new byte[1000];
        while (input.available() > 0) {
            int length = input.read(data);
        }*/


        if(out.contains("ConnectionClosed")) {
            socket.close();
            return out;
        }
        String msg = "";
        try{
            BufferedReader br = new BufferedReader(new InputStreamReader(input, "utf-8"));

            while(true)
                msg = msg.concat(br.readLine());
        }
        catch (Exception e)
        {

        }
        if(msg.contains("221"))
            out = out+"Closed";
        //System.out.println(new String(msg, "UTF-8"));
        return out;
    }

    void sendMessage(byte[] msg) throws Exception {
        if(!socket.isOutputShutdown())
            output.write(msg);
    }

    public void close() {
        if(targetProcess != null) {
            targetProcess.destroy();
        }
    }

    public String processSymbol(String input) throws Exception {
        String inAction = input;
        byte[] out = null;

        if(!socket.isConnected() || socket.isClosed()) return "ConnectionClosed";

        try {
            out = mapper.get(inAction);
            sendMessage(out);
            currentCase.add(out);
            return receiveMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";
            return outAction;
        }
    }

    public String sendFuzzingMessage (byte[] message) throws Exception
    {
        try {
            if (socket.isClosed())
                return "ConnectionClosed";
            output.write(message);
            return receiveFuzzingMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }
    }

    public String receiveFuzzingMessages() throws Exception {
        try{
            String out = receiveMessages();
            if(out.contains("ConnectionClosed")) {
                socket.close();
            }
            return out;
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }

    }

    //change by pany
    public void seedSave(List<byte[]> seed, String filename) {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        try {
            file = new File(filename);
            if (file.exists()) {
                return;
            }
            fos = new FileOutputStream(file);
            bos = new BufferedOutputStream(fos);
            for (int i = 0; i < seed.size(); i++) {
                byte[] tmp = intToByteArray(seed.get(i).length);
                bos.write(tmp, 0, 4);
                bos.write(seed.get(i), 0, seed.get(i).length);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (bos != null) {
                try {
                    bos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
    }
}

